"""
Copyright (c) 2024 [XKW.Beijing]
All rights reserved.

Author: [tangxiaojun]
Email: [417281862@qq.com]

LTP 是一个中文自然语言处理工具包
LTP（Language Technology Platform）是一个中文自然语言处理工具包，由中国科学院计算技术研究所（Institute of Computing Technology, Chinese Academy of Sciences）的自然语言处理研究团队开源。
这个工具包提供了多种文本处理功能，包括词性标注、句法分析、命名实体识别等，广泛应用于中文信息处理领域。

该Model需要python3.12版本运行，相应的包也需要该版本
分词 cws、
词性 pos、
命名实体标注 ner、
语义角色标注 srl：我们可以自动地识别出句子中的这些成分及其对应的语义角色。这在信息提取、文本理解、自动文摘、对话系统等领域有着广泛的应用。
依存句法分析 dep、
语义依存分析树 sdp、
语义依存分析图 sdpg
"""
import torch
from ltp import LTP

class LtpResource:

    def __init__(self):
        self.ltp = LTP("LTP/small")

    def ner_ltp(self, output, type):
        # 初始化一个空列表用于存储提取到的指定类型实体
        _names = []
        # 遍历分词结果及其对应的命名实体标签
        for word, tag in zip(output.cws, output.ner):
            _names = [name for label, name in tag if label == type]
            print(word,tag)

        # 返回提取到的指定类型实体列表
        print(type, _names)
        return _names

    # 预测文本分类，text需要删除标点
    def predict_model_by_text(self, text):
        output = self.ltp.pipeline([text], tasks=["cws", "pos", "ner", "srl", "dep", "sdp", "sdpg"])
        dict={}
        nh = self.ner_ltp(output, 'Nh')
        dict['姓名'] = nh
        ns = self.ner_ltp(output, 'Ns')
        dict['地址'] = ns
        return dict

    def predict_words_ltp(self,text):
        output = self.ltp.pipeline([text], tasks=["cws", "pos", "ner", "srl", "dep", "sdp", "sdpg"])
        dict={}
        dict['cws'] = output.cws
        dict['pos'] = output.pos
        dict['ner'] = output.ner
        dict['srl'] = output.srl
        dict['dep'] = output.dep
        dict['sdp'] = output.sdp
        dict['sdpg'] = output.sdpg
        return dict
